import numpy as np
import heapq
import math
import random
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d import Axes3D
from collections import deque
import matplotlib.animation as animation

# -----------------------------
# 参数设置
# -----------------------------
CELL_SIZE = 1.0            # 元胞边长（米）
FLOORS = 10                # 楼层数
HEIGHT, WIDTH = 30, 30     # 每层网格尺寸
AGENTS_PER_FLOOR = 40      # 每层人数
SPEED_MEAN = 1.0           # 初始速度均值（m/s）
SPEED_STD = 0.1            # 速度标准差
STAIR_SPEED_FACTOR = 0.7   # 楼梯速度系数
FLOOR_HEIGHT = 3.0         # 楼层高度（米）
MAX_STEPS = 1000           # 最大模拟步数
STAIR_COUNT = 3            # 每层楼梯数量
OBSTACLE_DENSITY = 0.25    # 迷宫墙体密度
ANIMATION_SPEED = 500      # 动画间隔（毫秒）

# 新增参数：出口拥堵度
MAX_QUEUE_LENGTH = 10      # 出口最大队列长度
CLOG_WEIGHT = 0.8          # 拥堵权重系数（0-1）

# -----------------------------
# 迷宫生成函数
# -----------------------------
def generate_maze(height, width):
    maze = np.zeros((height, width), dtype=int)
    for y in range(height):
        for x in range(width):
            if random.random() < OBSTACLE_DENSITY:
                maze[y, x] = 1
    start = (0, 0)
    end = (height-1, width-1)
    visited = np.zeros_like(maze, dtype=bool)
    stack = [start]
    while stack:
        y, x = stack[-1]
        if (y, x) == end:
            break
        neighbors = []
        for dy, dx in [(-1,0), (1,0), (0,-1), (0,1)]:
            ny, nx = y+dy, x+dx
            if 0<=ny<height and 0<=nx<width and not visited[ny, nx] and maze[ny, nx]==0:
                neighbors.append((ny, nx))
        if neighbors:
            ny, nx = random.choice(neighbors)
            maze[ny, nx] = 0
            visited[ny, nx] = True
            stack.append((ny, nx))
        else:
            stack.pop()
    return maze

# -----------------------------
# 环境构建函数
# -----------------------------
def build_environment():
    env = np.zeros((FLOORS, HEIGHT, WIDTH), dtype=int)
    exits = []
    
    # 生成各层迷宫
    for z in range(FLOORS):
        maze = generate_maze(HEIGHT, WIDTH)
        env[z] = maze
    
    # 设置楼梯和出口
    for z in range(FLOORS):
        for _ in range(STAIR_COUNT):
            while True:
                y = random.randint(0, HEIGHT-1)
                x = random.randint(0, WIDTH-1)
                if env[z, y, x] == 0:
                    env[z, y, x] = 2
                    break
    
    # 在顶层设置出口
    count = 0
    while count < STAIR_COUNT:
        y = random.randint(0, HEIGHT-1)
        x = random.randint(0, WIDTH-1)
        if env[FLOORS-1, y, x] == 2:
            env[FLOORS-1, y, x] = 3
            exits.append((FLOORS-1, y, x))
            count += 1
    return env, exits

# -----------------------------
# 出口类
# -----------------------------
class Exit:
    def __init__(self, position):
        self.position = position
        self.queue_length = 0
        self.clog_weight = 1.0
    
    def update_clog(self):
        self.clog_weight = 1.0 + (self.queue_length / MAX_QUEUE_LENGTH) * CLOG_WEIGHT
        self.clog_weight = min(self.clog_weight, 2.0)

# -----------------------------
# 动态静态场计算
# -----------------------------
def compute_static_field(env, exit_pos):
    f, h, w = env.shape
    static = np.full((f, h, w), np.inf)
    visited = np.zeros((f, h, w), dtype=bool)
    pq = []
    
    z0, y0, x0 = exit_pos
    static[z0, y0, x0] = 0.0
    heapq.heappush(pq, (0.0, z0, y0, x0))
    
    dirs = [(dz, dy, dx) for dz in (-1,0,1) 
            for dy in (-1,0,1) 
            for dx in (-1,0,1) 
            if not (dz==dy==dx==0)]
    
    while pq:
        dist_curr, z, y, x = heapq.heappop(pq)
        if visited[z, y, x]:
            continue
        visited[z, y, x] = True
        
        for dz, dy, dx in dirs:
            nz, ny, nx = z + dz, y + dy, x + dx
            if (0 <= nz < f and 0 <= ny < h and 0 <= nx < w and 
                env[nz, ny, nx] != 1):
                step = (CELL_SIZE * math.sqrt(1 + (FLOOR_HEIGHT/CELL_SIZE)**2) 
                        if dz != 0 else CELL_SIZE)
                nd = dist_curr + step
                if nd < static[nz, ny, nx]:
                    static[nz, ny, nx] = nd
                    heapq.heappush(pq, (nd, nz, ny, nx))
    return static

def compute_dynamic_static_field(env, exits):
    f, h, w = env.shape
    num_exits = len(exits)
    static = np.full((f, h, w, num_exits), np.inf)
    for i in range(num_exits):
        static[..., i] = compute_static_field(env, exits[i].position)
    return static

# -----------------------------
# A*算法实现
# -----------------------------
def a_star(static, exit_index, start, goal):
    open_set = []
    heapq.heappush(open_set, (0, start))
    came_from = {}
    g_score = {start: 0}
    f_score = {start: static[start[0], start[1], start[2], exit_index]}
    
    while open_set:
        current = heapq.heappop(open_set)[1]
        if current == goal:
            return reconstruct_path(came_from, current)
        
        for dz in (-1,0,1):
            for dy in (-1,0,1):
                for dx in (-1,0,1):
                    if dz == dy == dx == 0:
                        continue
                    neighbor = (current[0]+dz, current[1]+dy, current[2]+dx)
                    if (0 <= neighbor[0] < FLOORS and 
                        0 <= neighbor[1] < HEIGHT and 
                        0 <= neighbor[2] < WIDTH):
                        tentative_g = g_score[current] + static[neighbor[0], neighbor[1], neighbor[2], exit_index]
                        if tentative_g < g_score.get(neighbor, np.inf):
                            came_from[neighbor] = current
                            g_score[neighbor] = tentative_g
                            f_score[neighbor] = tentative_g + static[neighbor[0], neighbor[1], neighbor[2], exit_index]
                            heapq.heappush(open_set, (f_score[neighbor], neighbor))
    return []

def reconstruct_path(came_from, current):
    total_path = [current]
    while current in came_from:
        current = came_from[current]
        total_path.append(current)
    return total_path[::-1]

# -----------------------------
# Agent类
# -----------------------------
class Agent:
    def __init__(self, z, y, x):
        self.pos = (z, y, x)
        self.V0 = random.gauss(SPEED_MEAN, SPEED_STD)
        self.d_sum = 0.0
        self.t_sum = 0.0
        self.count = 0.0
        self.path = []
        self.current_exit = 0  # 当前选择的出口索引
        
    def update_path(self, static, exits):
        min_cost = np.inf
        best_exit = 0
        for i in range(len(exits)):
            cost = static[self.pos[0], self.pos[1], self.pos[2], i] * exits[i].clog_weight
            if cost < min_cost:
                min_cost = cost
                best_exit = i
        self.current_exit = best_exit  # 更新出口索引
        self.path = a_star(static, best_exit, self.pos, exits[best_exit].position)
        
    def speed(self, dz):
        base = self.V0 * (1 - 0.01 * (self.d_sum**0.71) / 
                   (1 - 0.01 * (self.t_sum**0.51) * 
                    math.exp(0.01*self.t_sum)))
        return base * (STAIR_SPEED_FACTOR if dz != 0 else 1.0)

# -----------------------------
# 模拟函数
# -----------------------------
def init_agents(env):
    agents = []
    for z in range(FLOORS):
        for _ in range(AGENTS_PER_FLOOR):
            while True:
                y = random.randint(0, HEIGHT-1)
                x = random.randint(0, WIDTH-1)
                if env[z,y,x] == 0:
                    agents.append(Agent(z, y, x))
                    env[z,y,x] = -1
                    break
    return agents

def simulate(env, static, exits):
    agents = init_agents(env.copy())
    frames = []
    stats = []
    dirs = [(dz, dy, dx) 
            for dz in (-1,0,1) 
            for dy in (-1,0,1) 
            for dx in (-1,0,1) 
            if not (dz==dy==dx==0)]
    
    for step in range(MAX_STEPS):
        occ = set(a.pos for a in agents)
        positions = []
        for a in agents[:]:
            a.t_sum += 1
            a.count += 1
            z, y, x = a.pos
            
            if step % 5 == 0:
                a.update_path(static, exits)
            
            next_pos = a.path[1] if len(a.path) > 1 else a.pos
            
            if next_pos in [exit.position for exit in exits]:
                for exit in exits:
                    if exit.position == next_pos:
                        exit.queue_length += 1
                agents.remove(a)
                continue
            
            # 显式访问静态场的当前出口维度
            current_exit = a.current_exit
            next_static = static[next_pos[0], next_pos[1], next_pos[2], current_exit]
            current_static = static[z, y, x, current_exit]
            
            if next_pos not in occ and next_static < current_static:
                new_pos = next_pos
                occ.remove(a.pos)
                occ.add(new_pos)
                a.pos = new_pos
                a.d_sum += CELL_SIZE
                a.count = 0
            else:
                a.count = 0
            positions.append(a.pos)
        
        for exit in exits:
            exit.update_clog()
        
        total = len(agents)
        stats.append((total, [exit.queue_length for exit in exits]))
        frames.append(positions)
        if not agents:
            break
    return frames, stats

# -----------------------------
# 可视化函数
# -----------------------------
def visualize(env, frames, stats, exits):
    cmap = ListedColormap(['white', 'dimgray', 'saddlebrown', 'green'])
    fig = plt.figure(figsize=(20, 10))
    gs = fig.add_gridspec(2, 6)
    
    ax3d = fig.add_subplot(gs[:, 0], projection='3d')
    ax3d.set_box_aspect([WIDTH, HEIGHT, FLOORS])
    ax3d.set_xlabel('X')
    ax3d.set_ylabel('Y')
    ax3d.set_zlabel('Floor')
    stairs = (env == 2)
    ax3d.voxels(stairs, facecolors='saddlebrown', edgecolor='k', alpha=0.7)
    exits_voxels = (env == 3)
    ax3d.voxels(exits_voxels, facecolors='green', edgecolor='k', alpha=0.7)
    scat3d = ax3d.scatter([], [], [], s=20, c='blue')
    text_total = ax3d.text2D(0.02, 0.95, '', transform=ax3d.transAxes)
    
    axs = []
    scats = []
    for f in range(FLOORS):
        row = 0 if f < 5 else 1
        col = (f % 5) + 1
        ax = fig.add_subplot(gs[row, col])
        ax.imshow(env[f], cmap=cmap, origin='lower')
        ax.set_title(f'Floor {f+1}')
        ax.axis('off')
        scats.append(ax.scatter([], [], s=10, c='red'))
        
    ax_clog = fig.add_subplot(gs[1, 5])
    ax_clog.set_title("Exit Clog Status")
    ax_clog.set_xlabel("Step")
    ax_clog.set_ylabel("Queue Length")
    
    def update(frame):
        pts = frames[frame]
        total, queues = stats[frame]
        xs, ys, zs = [], [], []
        for (z, y, x) in pts:
            xs.append(x)
            ys.append(y)
            zs.append(z)
        scat3d._offsets3d = (xs, ys, zs)
        text_total.set_text(f'总剩余: {total}')
        
        for f in range(FLOORS):
            floor_agents = [(y, x) for (z, y, x) in pts if z == f]
            if floor_agents:
                ys_f, xs_f = zip(*floor_agents)
                scats[f].set_offsets(np.column_stack([xs_f, ys_f]))
            else:
                scats[f].set_offsets(np.empty((0, 2)))
        
        ax_clog.clear()
        ax_clog.set_title("Exit Clog Status")
        ax_clog.set_xlabel("Step")
        ax_clog.set_ylabel("Queue Length")
        for i in range(len(exits)):
            ax_clog.plot([s[1][i] for s in stats[:frame+1]], label=f"Exit {i+1}")
        ax_clog.legend()
        
        return [scat3d] + scats + [text_total]
    
    anim = animation.FuncAnimation(
        fig, update, frames=len(frames), interval=ANIMATION_SPEED, blit=False
    )
    plt.show()

# -----------------------------
# 主程序
# -----------------------------
def main():
    env, exit_positions = build_environment()
    exits = [Exit(pos) for pos in exit_positions]
    static = compute_dynamic_static_field(env, exits)
    frames, stats = simulate(env, static, exits)
    visualize(env, frames, stats, exits)

if __name__ == '__main__':
    main()